import os
import sys
import subprocess
import platform

def setup():
    """设置财务数据可视化分析系统的运行环境"""
    print("=== 财务数据可视化分析系统 - 环境设置 ===")
    
    # 获取项目根目录
    root_dir = os.path.dirname(os.path.abspath(__file__))
    requirements_path = os.path.join(root_dir, "requirements.txt")
    
    # 检查requirements.txt是否存在
    if not os.path.exists(requirements_path):
        print(f"错误: 找不到依赖文件 {requirements_path}")
        return False
    
    # 创建必要的目录结构
    print("\n1. 创建必要的目录结构...")
    directories = [
        os.path.join(root_dir, "data", "raw"),
        os.path.join(root_dir, "data", "processed"),
    ]
    
    for directory in directories:
        if not os.path.exists(directory):
            try:
                os.makedirs(directory)
                print(f"  ✓ 创建目录: {directory}")
            except Exception as e:
                print(f"  ✗ 创建目录失败 {directory}: {e}")
                return False
        else:
            print(f"  ✓ 目录已存在: {directory}")
    
    # 询问是否创建虚拟环境
    print("\n2. 虚拟环境设置")
    create_venv = input("  是否创建虚拟环境? (推荐) [y/n]: ")
    
    if create_venv.lower() == 'y':
        venv_dir = os.path.join(root_dir, ".venv")
        
        if os.path.exists(venv_dir):
            print(f"  虚拟环境已存在: {venv_dir}")
            recreate = input("  是否重新创建? [y/n]: ")
            if recreate.lower() == 'y':
                try:
                    import shutil
                    shutil.rmtree(venv_dir)
                    print("  ✓ 已删除旧虚拟环境")
                except Exception as e:
                    print(f"  ✗ 删除旧虚拟环境失败: {e}")
                    return False
            else:
                print("  使用现有虚拟环境")
        
        if not os.path.exists(venv_dir):
            print("  正在创建虚拟环境...")
            try:
                subprocess.check_call([sys.executable, "-m", "venv", venv_dir])
                print(f"  ✓ 虚拟环境创建成功: {venv_dir}")
            except Exception as e:
                print(f"  ✗ 创建虚拟环境失败: {e}")
                return False
        
        # 激活虚拟环境并安装依赖
        print("\n3. 安装依赖")
        print("  正在安装依赖库...")
        
        # 根据操作系统确定激活脚本路径
        if platform.system() == "Windows":
            python_path = os.path.join(venv_dir, "Scripts", "python.exe")
        else:  # Linux or MacOS
            python_path = os.path.join(venv_dir, "bin", "python")
        
        if not os.path.exists(python_path):
            print(f"  ✗ 找不到虚拟环境Python解释器: {python_path}")
            return False
        
        try:
            subprocess.check_call([python_path, "-m", "pip", "install", "-r", requirements_path])
            print("  ✓ 依赖安装成功")
        except Exception as e:
            print(f"  ✗ 安装依赖失败: {e}")
            return False
        
        # 显示激活虚拟环境的命令
        print("\n4. 如何激活虚拟环境")
        if platform.system() == "Windows":
            print("  在命令提示符或PowerShell中运行:")
            print(f"  {venv_dir}\Scripts\activate")
        else:  # Linux or MacOS
            print("  在终端中运行:")
            print(f"  source {venv_dir}/bin/activate")
    else:
        # 直接安装依赖到系统Python
        print("\n3. 安装依赖")
        print("  正在安装依赖库到系统Python...")
        try:
            subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", requirements_path])
            print("  ✓ 依赖安装成功")
        except Exception as e:
            print(f"  ✗ 安装依赖失败: {e}")
            return False
    
    print("\n✅ 环境设置完成!")
    print("\n运行应用:")
    print("1. 使用Python直接运行: python run.py")
    print("2. 或者进入app目录运行: cd app && streamlit run app.py")
    
    return True

if __name__ == "__main__":
    setup()